// Copyright (c) 2020 Presto Labs Pte. Ltd.
// Author: hyuan

#include "coin2/exchange/okex_common/api_util/rate_tracker.h"

#include <algorithm>
#include <sstream>
#include <string>

#include <boost/algorithm/string/replace.hpp>
#include <glog/logging.h>

#include "coin2/base/proto_util.h"
#include "coin2/exchange/base/order/rate_config.h"

namespace coin2::exchange::okex_common::api_util {

using FixedWindowRateTracker = coin2::exchange::base::order::FixedWindowRateTracker;

constexpr auto EXAMPLE_SPOT_OKEX = R"(
"connection_config": {
  "key_filepath": "/home/ec2-user/config/trade/xunke18/okex/trade_key.json",
  "rate_limit_info": {
    "Spot.Okex.v5": {
      "ord_sec": {
        "secs": 2,
        "reqs": 60
      },
      "req_10_2sec": {
        "secs": 2,
        "reqs": 10
      },
      "req_20_2sec": {
        "secs": 2,
        "reqs": 20
      },
      "req_40_2sec": {
        "secs": 2,
        "reqs": 40
      },
      "req_60_2sec": {
        "secs": 2,
        "reqs": 60
      },
      "req_6_1sec": {
        "secs": 1,
        "reqs": 6
      }
    }
  }
},)";

template <class RateTracker>
void CheckFail(RateTracker&& rate_tracker, bool check_result) {
  if (!check_result) {
    LOG(INFO) << "Hit rate limit on [" << rate_tracker->GetName() << "]";
  }
}

OkexRateTracker::OkexRateTracker(const MarketExchangeApi& mea,
    const coin2::exchange::order::OrderConnectionConfig& cfg) {
  using RateConfig = coin2::exchange::base::order::RateConfig;
  if ((mea.market == coin::proto::MarketType::Spot) ||
      (mea.market == coin::proto::MarketType::Futures)) {
    coin2::exchange::order::RateLimitInfoSpotOkex rate_config;
    coin2::exchange::order::OrderConnectionConfig conn_cfg;
    conn_cfg.MergeFrom(cfg);
    RateConfig::Get()->GetConfig(mea, conn_cfg, &rate_config);
    if (mea.api == "v5" || mea.api == "v5-swap") {
      InitRateTrackerWithSecondTracker(&rate_config, mea.String(), EXAMPLE_SPOT_OKEX);
    } else {
      CHECK(false) << "unsupported mea: " << mea.String();
    }
  }
}

void OkexRateTracker::InitRateTrackerWithSecondTracker(const RateLimitInfoSpotOkex* rate_config,
    std::string_view mea_str, std::string_view example) {
  if (!rate_config || !rate_config->has_ord_sec() ||
      !rate_config->has_req_10_2sec() || !rate_config->has_req_20_2sec() ||
      !rate_config->has_req_40_2sec() || !rate_config->has_req_60_2sec() ||
      !rate_config->has_req_6_1sec()) {
    coin2::exchange::base::order::RateConfig::CheckFail(mea_str, example);
  }

  const auto& ord_sec = rate_config->ord_sec();
  const auto& req_10_2sec = rate_config->req_10_2sec();
  const auto& req_20_2sec = rate_config->req_20_2sec();
  const auto& req_40_2sec = rate_config->req_40_2sec();
  const auto& req_60_2sec = rate_config->req_60_2sec();
  const auto& req_6_1sec = rate_config->req_6_1sec();
  LOG(INFO) << "\n~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~"
      << "\nOG request rate was set as below:\n"
      << "\nord_sec: " << ord_sec.reqs() << "(reqs)/" << ord_sec.secs() << "(seconds)]"
      << "\nreq_10_2sec: " << req_10_2sec.reqs() << "(reqs) / " << req_10_2sec.secs() << "(seconds)]"
      << "\nreq_20_2sec: " << req_20_2sec.reqs() << "(reqs) / " << req_20_2sec.secs() << "(seconds)]"
      << "\nreq_40_2sec: " << req_40_2sec.reqs() << "(reqs) / " << req_40_2sec.secs() << "(seconds)]"
      << "\nreq_60_2sec: " << req_60_2sec.reqs() << "(reqs) / " << req_60_2sec.secs() << "(seconds)]"
      << "\nreq_6_1sec: " << req_6_1sec.reqs() << "(reqs) / " << req_6_1sec.secs() << "(seconds)]"
      << "\n\nOr you can configure customized rate in driver config. For example: " << example
      << "\n~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~";

  ord_reqs_ = ord_sec.reqs();
  ord_nsecs_ = ord_sec.secs() * 1'000'000'000LL;

  req_balance_rate_tracker_ =
      std::make_unique<FixedWindowRateTracker>(req_10_2sec.secs() * 1'000'000'000LL);
  req_balance_rate_tracker_->SetRateLimitInfo(req_10_2sec.reqs(), req_10_2sec.reqs());
  req_balance_rate_tracker_->SetName("req_balance_rate_tracker_(10/2sec)");

  req_position_rate_tracker_ =
      std::make_unique<FixedWindowRateTracker>(req_10_2sec.secs() * 1'000'000'000LL);
  req_position_rate_tracker_->SetRateLimitInfo(req_10_2sec.reqs(), req_10_2sec.reqs());
  req_position_rate_tracker_->SetName("req_position_rate_tracker_(10/2sec)");

  req_open_order_rate_tracker_ =
      std::make_unique<FixedWindowRateTracker>(req_20_2sec.secs() * 1'000'000'000LL);
  req_open_order_rate_tracker_->SetRateLimitInfo(req_20_2sec.reqs(), req_20_2sec.reqs());
  req_open_order_rate_tracker_->SetName("req_open_order_rate_tracker_(20/2sec)");

  req_hist_order_rate_tracker_ =
      std::make_unique<FixedWindowRateTracker>(req_40_2sec.secs() * 1'000'000'000LL);
  req_hist_order_rate_tracker_->SetRateLimitInfo(req_40_2sec.reqs(), req_40_2sec.reqs());
  req_hist_order_rate_tracker_->SetName("req_hist_order_rate_tracker_(40/2sec)");

  req_fill_order_rate_tracker_ =
      std::make_unique<FixedWindowRateTracker>(req_60_2sec.secs() * 1'000'000'000LL);
  req_fill_order_rate_tracker_->SetRateLimitInfo(req_60_2sec.reqs(), req_60_2sec.reqs());
  req_fill_order_rate_tracker_->SetName("req_fill_order_rate_tracker_(60/2sec)");

  req_deposit_rate_tracker_ =
      std::make_unique<FixedWindowRateTracker>(req_6_1sec.secs() * 1'000'000'000LL);
  req_deposit_rate_tracker_->SetRateLimitInfo(req_6_1sec.reqs(), req_6_1sec.reqs());
  req_deposit_rate_tracker_->SetName("req_deposit_rate_tracker_(6/1sec)");

  req_withdrawal_rate_tracker_ =
      std::make_unique<FixedWindowRateTracker>(req_6_1sec.secs() * 1'000'000'000LL);
  req_withdrawal_rate_tracker_->SetRateLimitInfo(req_6_1sec.reqs(), req_6_1sec.reqs());
  req_withdrawal_rate_tracker_->SetName("req_withdrawal_rate_tracker_(6/1sec)");

  rate_limit_report_writer_ = std::make_unique<RateLimitReportWriter>();
  rate_limit_report_writer_->SetRateLimitInfoString(rate_config);
}

FixedWindowRateTracker* OkexRateTracker::GetOrderRateTracker(
    OrderRateTrackerMap* ord_rate_tracker_map,
    const std::string& map_name,
    const std::string& instrument) {
  auto itr = ord_rate_tracker_map->find(instrument);
  if (itr != ord_rate_tracker_map->end()) {
    return itr->second.get();
  } else {
    std::unique_ptr<FixedWindowRateTracker> rate_tracker =
        std::make_unique<FixedWindowRateTracker>(ord_nsecs_);
    rate_tracker->SetRateLimitInfo(ord_reqs_, ord_reqs_);
    rate_tracker->SetName(map_name + "[" + instrument + "](60/2sec)");
    ord_rate_tracker_map->insert(std::make_pair(instrument, std::move(rate_tracker)));
    return (*ord_rate_tracker_map)[instrument].get();
  }
}

FixedWindowRateTracker* OkexRateTracker::GetPlaceOrderRateTracker(
    const std::string& instrument) {
  return GetOrderRateTracker(&ord_place_order_rate_tracker_map_,
      "ord_place_order_rate_tracker", instrument);
}

FixedWindowRateTracker* OkexRateTracker::GetCancelOrderRateTracker(
    const std::string& instrument) {
  return GetOrderRateTracker(&ord_cancel_order_rate_tracker_map_,
      "ord_cancel_order_rate_tracker", instrument);
}

bool OkexRateTracker::HasEnoughRateLimit(int64_t timestamp, ActionType action_type,
    QueryType query_type, const std::string& instrument) {
  bool ret = true;
  if (action_type == ActionType::PLACE) {
    FixedWindowRateTracker* place_order_rate_tracker = GetPlaceOrderRateTracker(instrument);
    if (ret && place_order_rate_tracker) {
      ret &= (place_order_rate_tracker->GetRemaining(timestamp) > 2);
      CheckFail(place_order_rate_tracker, ret);
    }
  } else if (action_type == ActionType::CANCEL) {
    FixedWindowRateTracker* cancel_order_rate_tracker = GetCancelOrderRateTracker(instrument);
    if (ret && cancel_order_rate_tracker) {
      ret &= (cancel_order_rate_tracker->GetRemaining(timestamp) > 2);
      CheckFail(cancel_order_rate_tracker, ret);
    }
  } else {
    if (query_type == QueryType::BALANCE) {
      if (ret && req_balance_rate_tracker_) {
        ret &= (req_balance_rate_tracker_->GetRemaining(timestamp) > 2);
        CheckFail(req_balance_rate_tracker_, ret);
      }
    } else if (query_type == QueryType::POSITION) {
      if (ret && req_position_rate_tracker_) {
        ret &= (req_position_rate_tracker_->GetRemaining(timestamp) > 2);
        CheckFail(req_position_rate_tracker_, ret);
      }
    } else if (query_type == QueryType::OPEN_ORDER) {
      if (ret && req_open_order_rate_tracker_) {
        ret &= (req_open_order_rate_tracker_->GetRemaining(timestamp) > 2);
        CheckFail(req_open_order_rate_tracker_, ret);
      }
    } else if (query_type == QueryType::HIST_ORDER) {
      if (ret && req_hist_order_rate_tracker_) {
        ret &= (req_hist_order_rate_tracker_->GetRemaining(timestamp) > 2);
        CheckFail(req_hist_order_rate_tracker_, ret);
      }
    } else if (query_type == QueryType::FILL_ORDER) {
      if (ret && req_fill_order_rate_tracker_) {
        ret &= (req_fill_order_rate_tracker_->GetRemaining(timestamp) > 2);
        CheckFail(req_fill_order_rate_tracker_, ret);
      }
    } else if (query_type == QueryType::DEPOSIT) {
      if (ret && req_deposit_rate_tracker_) {
        ret &= (req_deposit_rate_tracker_->GetRemaining(timestamp) > 2);
        CheckFail(req_deposit_rate_tracker_, ret);
      }
    } else if (query_type == QueryType::WITHDRAWAL) {
      if (ret && req_withdrawal_rate_tracker_) {
        ret &= (req_withdrawal_rate_tracker_->GetRemaining(timestamp) > 2);
        CheckFail(req_withdrawal_rate_tracker_, ret);
      }
    }
  }
  if (!ret) {
    rate_limit_report_writer_->SetHitRateLimit(timestamp, action_type);
  }
  return ret;
}

void OkexRateTracker::Push(int32_t num_reqs, int64_t timestamp, ActionType action_type,
    QueryType query_type, const std::string& instrument) {
  if (action_type == ActionType::PLACE) {
    FixedWindowRateTracker* place_order_rate_tracker = GetPlaceOrderRateTracker(instrument);
    if (place_order_rate_tracker) {
      place_order_rate_tracker->Push(num_reqs, timestamp);
    }
  } else if (action_type == ActionType::CANCEL) {
    FixedWindowRateTracker* cancel_order_rate_tracker = GetCancelOrderRateTracker(instrument);
    if (cancel_order_rate_tracker) {
      cancel_order_rate_tracker->Push(num_reqs, timestamp);
    }
  } else {
    if (query_type == QueryType::BALANCE) {
      if (req_balance_rate_tracker_) {
        req_balance_rate_tracker_->Push(num_reqs, timestamp);
      }
    } else if (query_type == QueryType::POSITION) {
      if (req_position_rate_tracker_) {
        req_position_rate_tracker_->Push(num_reqs, timestamp);
      }
    } else if (query_type == QueryType::OPEN_ORDER) {
      if (req_open_order_rate_tracker_) {
        req_open_order_rate_tracker_->Push(num_reqs, timestamp);
      }
    } else if (query_type == QueryType::HIST_ORDER) {
      if (req_hist_order_rate_tracker_) {
        req_hist_order_rate_tracker_->Push(num_reqs, timestamp);
      }
    } else if (query_type == QueryType::FILL_ORDER) {
      if (req_fill_order_rate_tracker_) {
        req_fill_order_rate_tracker_->Push(num_reqs, timestamp);
      }
    } else if (query_type == QueryType::DEPOSIT) {
      if (req_deposit_rate_tracker_) {
        req_deposit_rate_tracker_->Push(num_reqs, timestamp);
      }
    } else if (query_type == QueryType::WITHDRAWAL) {
      if (req_withdrawal_rate_tracker_) {
        req_withdrawal_rate_tracker_->Push(num_reqs, timestamp);
      }
    }
  }
}

void OkexRateTracker::SetRateLimitInfoString(const google::protobuf::Message* message) {
  if (!rate_limit_report_writer_) {
    return;
  }
  rate_limit_report_writer_->SetRateLimitInfoString(message);
}

void OkexRateTracker::WriteRateLimitReportMessage(OrderLoggerCB callback) {
  if (!callback) {
    return;
  }
  rate_limit_report_writer_->Write(callback);
}

double OkexRateTracker::GetMaxRateUsedRatio(int64_t timestamp, ActionType action_type) {
  switch (action_type) {
    case ActionType::PLACE: {
      return GetMaxRateUsedRatio(ord_place_order_rate_tracker_map_, timestamp);
    }
    case ActionType::CANCEL: {
      return GetMaxRateUsedRatio(ord_cancel_order_rate_tracker_map_, timestamp);
    }
    case ActionType::QUERY: {
      double balance_rate_used_ratios =
          GetMaxRateUsedRatio(req_balance_rate_tracker_, timestamp);
      double position_rate_used_ratios =
          GetMaxRateUsedRatio(req_position_rate_tracker_, timestamp);
      double open_order_rate_used_ratios =
          GetMaxRateUsedRatio(req_open_order_rate_tracker_, timestamp);
      double hist_order_rate_used_ratios =
          GetMaxRateUsedRatio(req_hist_order_rate_tracker_, timestamp);
      double fill_order_rate_used_ratios =
          GetMaxRateUsedRatio(req_fill_order_rate_tracker_, timestamp);
      double deposit_rate_used_ratios =
          GetMaxRateUsedRatio(req_deposit_rate_tracker_, timestamp);
      double withdrawal_rate_used_ratios =
          GetMaxRateUsedRatio(req_withdrawal_rate_tracker_, timestamp);
      return std::max(
          std::max(std::max(balance_rate_used_ratios, position_rate_used_ratios),
              std::max(open_order_rate_used_ratios, hist_order_rate_used_ratios)
          ),
          std::max(fill_order_rate_used_ratios,
              std::max(deposit_rate_used_ratios, withdrawal_rate_used_ratios)
          ));
    }
    default: {
      LOG(ERROR) << "ops~ something is wrong, action_type=" << action_type;
      return 0.;
    }
  }
}

}  // namespace coin2::exchange::okex_common::api_util
